import time
import pyautogui as pa
import pywinauto as pw
import schedule
import yagmail
import requests
import json
import random
from datetime import datetime, timedelta
import logging
import psutil
import os
from typing import List, Optional, Dict, Any
import threading

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('qmt_auto_login.log', encoding='utf-8'),
        logging.StreamHandler()
    ]
)
class QMTAutoLogin:
    '''
    QMT自动登录系统 - 支持多时间点登录和退出
    作者:索普量化
    微信:xms_quants1
    '''
    def __init__(self, 
                 connect_path: str = r'D:\国金QMT交易端模拟\\bin.x64\XtItClient.exe',
                 user: str = '', 
                 password: str = '', 
                 notify_type: str = 'qq',
                 email_qq: str = '1752515969@qq.com', 
                 email_password: str = 'jgyhavzupyglecaf',
                 access_token_list: List[str] = ['1029762153@qq.com'],
                 login_times: List[str] = ["09:00", "13:00"],  # 多个登录时间
                 logout_times: List[str] = ["11:30", "15:00"], # 多个退出时间
                 trading_days: List[int] = [0, 1, 2, 3, 4]):  # 0=周一, 4=周五
        '''
        Args:
            connect_path: qmt安装路径
            user: 股票账户
            password: 密码
            notify_type: 通知方式 wx微信,qq qq,dd钉钉
            email_qq: QQ邮箱地址
            email_password: QQ邮箱授权码
            access_token_list: 通知token列表
            login_times: 自动登录时间列表，格式 ["HH:MM", "HH:MM"]
            logout_times: 自动退出时间列表，格式 ["HH:MM", "HH:MM"]
            trading_days: 交易日期列表 [0,1,2,3,4] 表示周一到周五
        '''
        self.connect_path = connect_path
        self.user = user
        self.password = password
        self.app = None
        self.notify_type = notify_type
        self.access_token_list = access_token_list
        self.email_qq = email_qq
        self.email_password = email_password
        self.login_times = login_times
        self.logout_times = logout_times
        self.trading_days = trading_days
        
        # 状态管理
        self.is_logged_in = False
        self.last_login_time = None
        self.login_attempts = 0
        self.max_login_attempts = 5
        self.login_retry_delay = 30  # 重试延迟(秒)
        self.last_attempt_time = None
        self.consecutive_failures = 0
        self.max_consecutive_failures = 3
        self.schedule_lock = threading.Lock()
        
        # 初始化任务列表
        self.scheduled_jobs = []
        
        # 设置定时任务
        self._setup_schedule()

    def _setup_schedule(self):
        """设置定时任务"""
        # 清除现有任务
        schedule.clear()
        self.scheduled_jobs.clear()
        
        # 为每个登录时间设置任务
        for login_time in self.login_times:
            job = schedule.every().day.at(login_time).do(self.scheduled_login, login_time)
            self.scheduled_jobs.append(job)
            logging.info(f"设置登录任务: 每天 {login_time}")
        
        # 为每个退出时间设置任务
        for logout_time in self.logout_times:
            job = schedule.every().day.at(logout_time).do(self.scheduled_logout, logout_time)
            self.scheduled_jobs.append(job)
            logging.info(f"设置退出任务: 每天 {logout_time}")
        
        # 每分钟检查连接状态
        schedule.every(1).minutes.do(self.run_log)
        
        # 每30分钟健康检查
        schedule.every(30).minutes.do(self.health_check)
        
        # 每天凌晨清理状态
        schedule.every().day.at("00:00").do(self.daily_cleanup)
        
        logging.info(f"定时任务设置完成 - 登录时间: {self.login_times}, 退出时间: {self.logout_times}")

    def is_trading_day(self) -> bool:
        """检查今天是否是交易日"""
        today = datetime.now().weekday()  # 0=周一, 6=周日
        return today in self.trading_days

    def should_perform_operation(self) -> bool:
        """检查是否应该执行操作（登录/退出）"""
        if not self.is_trading_day():
            logging.info("今天不是交易日，跳过操作")
            return False
        return True

    def send_dingding(self, msg: str = '买卖交易成功,', access_token_list: Optional[List[str]] = None):
        '''发送钉钉通知'''
        if access_token_list is None:
            access_token_list = self.access_token_list
            
        access_token = random.choice(access_token_list)
        url = f'https://oapi.dingtalk.com/robot/send?access_token={access_token}'
        headers = {'Content-Type': 'application/json;charset=utf-8'}
        data = {
            "msgtype": "text",
            "at": {"isAtAll": False},
            "text": {"content": msg}
        }
        try:
            r = requests.post(url, data=json.dumps(data), headers=headers, timeout=10)
            text = r.json()
            if text['errmsg'] == 'ok':
                logging.info('钉钉发送成功')
                return text
            else:
                logging.error(f'钉钉发送失败: {text}')
                return text
        except Exception as e:
            logging.error(f'钉钉发送异常: {e}')
            return None

    def send_email(self, text: str = '交易完成', email_qq: Optional[str] = None, 
                   email_password: Optional[str] = None, receiver: Optional[str] = None):
        '''发送QQ邮件'''
        if email_qq is None:
            email_qq = self.email_qq
        if email_password is None:
            email_password = self.email_password
        if receiver is None:
            receiver = self.access_token_list[-1]
            
        try:
            yag = yagmail.SMTP(user=email_qq, password=email_password, host='smtp.qq.com')
            yag.send(to=receiver, contents=text, subject='QMT系统通知')
            logging.info('QQ邮件发送完成')
            return True
        except Exception as e:
            logging.error(f'QQ邮件发送失败: {e}')
            return False

    def send_wechat(self, msg: str = '买卖交易成功,', access_token_list: Optional[List[str]] = None):
        '''发送微信通知'''
        if access_token_list is None:
            access_token_list = self.access_token_list
            
        access_token = random.choice(access_token_list)
        url = f'https://qyapi.weixin.qq.com/cgi-bin/webhook/send?key={access_token}'
        headers = {'Content-Type': 'application/json;charset=utf-8'}
        data = {
            "msgtype": "text",
            "at": {"isAtAll": False},
            "text": {"content": msg}
        }
        try:
            r = requests.post(url, data=json.dumps(data), headers=headers, timeout=10)
            text = r.json()
            if text['errmsg'] == 'ok':
                logging.info('微信发送成功')
                return text
            else:
                logging.error(f'微信发送失败: {text}')
                return text
        except Exception as e:
            logging.error(f'微信发送异常: {e}')
            return None

    def send_notification(self, text: str = ''):
        '''发送通知信息'''
        logging.info(f'发送消息: {text}')
        
        if self.notify_type == 'qq':
            self.send_email(text=text)
        elif self.notify_type == 'wx':
            self.send_wechat(msg=text)
        elif self.notify_type == 'dd':
            self.send_dingding(msg=text)
        else:
            self.send_email(text=text)

    def is_process_running(self, process_name: str = "XtItClient.exe") -> bool:
        """检查QMT进程是否在运行"""
        for proc in psutil.process_iter(['name']):
            if proc.info['name'] == process_name:
                return True
        return False

    def verify_login_success(self, timeout: int = 15) -> bool:
        """验证登录是否成功"""
        start_time = time.time()
        while time.time() - start_time < timeout:
            try:
                # 检查进程是否运行
                if not self.is_process_running():
                    logging.warning("QMT进程未运行")
                    return False
                
                # 尝试查找QMT主窗口
                test_app = pw.application.Application(backend="uia")
                proc_id = pw.application.process_from_module("XtItClient.exe")
                app = test_app.connect(process=proc_id, timeout=5)
                
                # 查找主窗口
                main_window = app.window_(title_re=".*QMT.*", control_type="Window")
                if main_window.exists():
                    logging.info("QMT登录验证成功 - 找到主窗口")
                    self.is_logged_in = True
                    self.last_login_time = datetime.now()
                    self.consecutive_failures = 0
                    return True
                    
            except Exception as e:
                logging.debug(f"登录验证尝试失败: {e}")
                
            time.sleep(2)
            
        logging.warning("QMT登录验证失败 - 未找到主窗口")
        return False

    def should_retry_login(self) -> bool:
        """判断是否应该重试登录"""
        # 如果已经登录，不需要重试
        if self.is_logged_in:
            return False
            
        # 如果没有尝试过登录，应该重试
        if self.last_attempt_time is None:
            return True
            
        # 检查连续失败次数
        if self.consecutive_failures >= self.max_consecutive_failures:
            logging.warning(f"连续登录失败次数已达{self.consecutive_failures}次，暂停重试")
            return False
            
        # 检查距离上次尝试的时间
        time_since_last_attempt = (datetime.now() - self.last_attempt_time).total_seconds()
        if time_since_last_attempt < self.login_retry_delay:
            return False
            
        return True

    def login_with_retry(self, max_attempts: int = 3) -> bool:
        """带重试的登录方法"""
        for attempt in range(max_attempts):
            try:
                logging.info(f"尝试登录QMT (第{attempt + 1}次)")
                
                # 启动新的QMT实例
                self.app = pw.Application(backend='uia').start(self.connect_path, timeout=10)
                time.sleep(10)  # 等待程序启动
                
                pa.hotkey('ctrl', 'a')
                time.sleep(1)
                # 删除选中的内容
                pa.press('backspace')
                # 输入用户名
                pa.typewrite(self.user)
                time.sleep(2)
                pa.hotkey('enter')
                pa.hotkey('ctrl', 'a')
                time.sleep(1)
                # 删除选中的内容
                pa.press('backspace')
                time.sleep(2)
                pa.typewrite(self.password)
                time.sleep(2)
                pa.hotkey('enter')
                time.sleep(1)
                pa.hotkey('enter')
                time.sleep(3)  # 等待登录完成
                
                # 验证登录是否成功
                if False:
                    logging.info("QMT登录成功")
                    self.send_notification(f"QMT自动登录成功 - {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
                    return True
                else:
                    logging.warning(f"第{attempt + 1}次登录验证失败")
                    if attempt < max_attempts - 1:
                        time.sleep(self.login_retry_delay)
                    
            except Exception as e:
                logging.error(f"第{attempt + 1}次登录尝试异常: {e}")
                if attempt < max_attempts - 1:
                    time.sleep(self.login_retry_delay)
        
        logging.error(f"经过{max_attempts}次尝试后登录失败")
        self.send_notification(f"QMT自动登录失败 - 请检查系统状态")
        return False

    def login(self) -> bool:
        """登录QMT（主登录方法）"""
        self.last_attempt_time = datetime.now()
        
        # 检查是否已经在运行
        #if self.is_process_running():
            #logging.info("QMT进程已在运行，检查登录状态")
            #if self.verify_login_success():
                #return True
            
        
        # 执行登录
        success = self.login_with_retry(max_attempts=3)
        if success:
            self.consecutive_failures = 0
        else:
            self.consecutive_failures += 1
            
        return success

    def scheduled_login(self, login_time: str):
        """定时登录任务"""
        with self.schedule_lock:
            logging.info(f"执行定时登录任务: {login_time}")
            
            if not self.should_perform_operation():
                return
                
            current_time = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
            self.send_notification(f"开始执行定时登录 {login_time} - {current_time}")
            
            # 如果已经登录，跳过
            if self.is_logged_in:
                logging.info("系统已登录，跳过登录操作")
                return
                
            self.login()

    def scheduled_logout(self, logout_time: str):
        """定时退出任务"""
        with self.schedule_lock:
            logging.info(f"执行定时退出任务: {logout_time}")
            
            if not self.should_perform_operation():
                return
                
            current_time = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
            self.send_notification(f"开始执行定时退出 {logout_time} - {current_time}")
            
            # 如果已经退出，跳过
            if not self.is_logged_in:
                logging.info("系统已退出，跳过退出操作")
                return
                
            self.kill()

    def kill(self):
        '''退出程序'''
        try:
            if self.app:
                self.app.kill()
            # 强制终止进程
            for proc in psutil.process_iter(['name']):
                if proc.info['name'] == "XtItClient.exe":
                    proc.terminate()
                    
            self.is_logged_in = False
            logging.info('QMT已退出')
            self.send_notification(f'{datetime.now()} QMT已安全退出')
        except Exception as e:
            logging.error(f'退出QMT时发生异常: {e}')

    def check_and_reconnect(self) -> bool:
        """检查连接状态并在需要时重新连接"""
        # 只在交易时间检查重连
        if not self.should_perform_operation():
            return self.is_logged_in
            
        current_status = self.verify_login_success(timeout=5)
        
        if current_status and not self.is_logged_in:
            logging.info("检测到QMT已连接但状态未更新，更新状态")
            self.is_logged_in = True
            self.last_login_time = datetime.now()
            self.consecutive_failures = 0
            return True
        elif not current_status and self.is_logged_in:
            logging.warning("检测到QMT连接断开，尝试重新连接")
            self.is_logged_in = False
            return self.login()
            
        return self.is_logged_in

    def auto_retry_login(self) -> bool:
        """自动重试登录"""
        if not self.should_retry_login():
            return False
            
        logging.info("自动重试登录QMT...")
        return self.login()

    def run_log(self):
        '''运行状态检查和维护'''
        current_time = datetime.now()
        
        # 只在交易时间检查连接状态
        if self.should_perform_operation():
            self.check_and_reconnect()
            
            # 如果未登录且应该重试，尝试登录
            if not self.is_logged_in and self.should_retry_login():
                logging.info("检测到QMT未登录，尝试登录...")
                self.login()
        
        # 记录状态日志
        status = "已登录" if self.is_logged_in else "未登录"
        consecutive_failures_info = f"，连续失败次数: {self.consecutive_failures}" if self.consecutive_failures > 0 else ""
        log_msg = f'系统运行中 - 状态: {status}{consecutive_failures_info}'
        
        # 每15分钟记录一次详细状态
        if current_time.minute % 15 == 0:
            logging.info(log_msg)

    def health_check(self):
        """健康检查"""
        logging.info("执行健康检查")
        if self.should_perform_operation() and not self.is_logged_in:
            self.send_notification("QMT健康检查: 系统未登录，尝试重新连接")
            self.auto_retry_login()
        else:
            logging.info("健康检查: 系统运行正常")

    def daily_cleanup(self):
        """每日清理任务"""
        logging.info("执行每日清理任务")
        self.consecutive_failures = 0
        self.login_attempts = 0

    def update_schedule(self, login_times: List[str] = None, logout_times: List[str] = None):
        """更新调度时间"""
        if login_times:
            self.login_times = login_times
        if logout_times:
            self.logout_times = logout_times
            
        self._setup_schedule()
        logging.info(f"已更新调度时间 - 登录: {self.login_times}, 退出: {self.logout_times}")

    def get_schedule_info(self) -> Dict[str, Any]:
        """获取调度信息"""
        return {
            "login_times": self.login_times,
            "logout_times": self.logout_times,
            "trading_days": self.trading_days,
            "is_logged_in": self.is_logged_in,
            "last_login_time": self.last_login_time,
            "consecutive_failures": self.consecutive_failures
        }

    def cleanup(self):
        '''清理资源'''
        if self.is_logged_in:
            self.kill()
        # 清理调度任务
        schedule.clear()

    def run(self):
        """运行主循环"""
        logging.info('QMT自动登录系统启动成功，开始监控...')
        logging.info(f'交易日期: {self.trading_days}, 登录时间: {self.login_times}, 退出时间: {self.logout_times}')
        
        # 启动时立即尝试登录一次（如果是交易日）
        if self.should_perform_operation():
            logging.info("启动时立即尝试登录")
            self.login()
        else:
            logging.info("今天不是交易日，跳过启动登录")
        
        try:
            while True:
                schedule.run_pending()
                time.sleep(1)
        except KeyboardInterrupt:
            logging.info("收到中断信号，正在关闭...")
        finally:
            self.cleanup()

if __name__ == '__main__':
    # 初始化QMT自动登录
    qmt_auto = QMTAutoLogin(
        connect_path=r'D:\国金QMT交易端模拟/bin.x64/XtItClient.exe',
        user='55001947',
        password='259800',
        notify_type='qq',
        email_qq='1752515969@qq.com',
        email_password='jgyhavzupyglecaf',
        access_token_list=['12153@qq.com'],
        login_times=["09:00", "22:00"],  # 上午9点和下午1点登录
        logout_times=["08:40", "21:00"], # 上午11:30和下午3点退出
        trading_days=[0, 1, 2, 3, 4]    # 周一到周五
    )
    #测试
    qmt_auto.login()
    # 运行系统
    qmt_auto.run()